#include "Calculus.h"

#include <cmath>
#include <climits>
#include <cfloat>
#include <cassert>

#define MIN(a, b) ((a) < (b) ? (a) : (b))
#define MAX(a, b) ((a) > (b) ? (a) : (b))
#define ACCURACY 1e-8
#define ITERATION 1000

using namespace Calculus;
using namespace LinearAlgebra;
using namespace UniPolynomial;

// 多元向量微分，返回 Jacobi 矩阵
Matrix Jacobian::operator()(FuncXX f, const Vector &x)
{
    Vector y = f(x);
    int m = y.size();
    int n = x.size();
    Real h = 1e-8;

    Matrix d(m, n);
    for (int i = 0; i < n; i++)
    {
        Vector y1 = x, y2 = x;
        y1[i] += h, y2[i] -= h;

        // 中心差分向量
        Vector dp = (f(y1) - f(y2)) / (2 * h);
        for (int j = 0; j < m; j++)
            d(j, i) = dp[j];
    }
    return d;
}

// 多元函数的 Jacobi 矩阵：由对每个分量的偏导构成的矩阵
Matrix Jacobian::operator()(FuncX1 f, const Vector &x)
{
    int n = x.size();
    Real h = 1e-8;

    // 1行矩阵
    Matrix d(1, n);
    for (int i = 0; i < n; i++)
    {
        Vector y1 = x, y2 = x;
        y1[i] += h, y2[i] -= h;

        // 中心差分计算导数
        d(0, i) = (f(y1) - f(y2)) / (2 * h);
    }
    return d;
}

// 一元向量函数的导向量：每个分量的导数构成的向量
Vector Jacobian::operator()(Func1X f, Real x)
{
    Real h = 1e-8;

    // 直接计算向量作差
    Vector y1 = f(x + h);
    Vector y2 = f(x - h);

    return (y1 - y2) / (2 * h);
}

// 一元函数的导数
Real Jacobian::operator()(Func11 f, Real x)
{
    Real h = 1e-8;
    return (f(x + h) - f(x - h)) / (2 * h);
}

// 初始化权
Integrator::Integrator(int N) : m_weight(Vector(N + 1))
{
}

// 积分器分割数
int Integrator::order() const
{
    return m_weight.size() - 1;
}

// 初始化权重
NewtonCotes::NewtonCotes(int N) : Integrator(N)
{
    // 使用插值多项式积分计算权
    Vector x(N + 1);
    for (int i = 0; i < N + 1; i++)
        x[i] = i;

    for (int i = 0; i < N + 1; i++)
    {
        Vector f(N + 1);
        f[i] = 1;
        m_weight[i] = interpolate(x, f).intergral(0, N);
    }
}

// 牛顿-柯茨公式
Real NewtonCotes::operator()(Func11 f, Real a, Real b)
{
    int N = m_weight.size() - 1;

    // 分割宽度
    Real h = (b - a) / N;

    // 节点值
    Vector y(N + 1);
    for (int i = 0; i < N + 1; i++)
        y[i] = f(h * i + a);

    // 节点加权
    return dot(m_weight, y) * h;
}

// 初始化权节点和权重
GaussIntegrator::GaussIntegrator(int N) : Integrator(N), m_node(Vector(N + 1))
{
    // 获得 Gauss-Legendre 多项式
    Polynomial GL = Legendre(N + 1);

    // 使用牛顿法迭代求根
    FSolve FS;

    // 依次计算 N+1 个根，从 0 开始迭代，通常会得到从小到大的根
    for (int i = 0; i < N + 1; i++)
    {
        // 构造函数
        Func11 F = [=](Real x)
        {
            return GL(x);
        };

        // 迭代求根，从 1 开始，因为 Legendre 多项式的根不在 -1,1 之外，防止初始值与根重合，或者出现导数为零的情况
        m_node[i] = FS(F, 1);

        // 除以因式 x - c[i] 进行降阶
        GL = GL / Polynomial({-m_node[i], 1});
    }

    // 将根按照升序排列
    m_node.sort();

    // 初始化权，在 -1,1 的正交多项式区间上对初等拉格朗日函数积分
    for (int i = 0; i < N + 1; i++)
        m_weight[i] = Lagrange(m_node, i).intergral(-1, 1);
}

// 高斯积分公式
Real GaussIntegrator::operator()(Func11 f, Real a, Real b)
{
    int N = m_weight.size() - 1;

    // 节点值，需要区间从 -1,1 变换到 a,b
    Vector y(N + 1);
    for (int i = 0; i < N + 1; i++)
        y[i] = (b - a) / 2 * f((b - a) / 2 * m_node[i] + (b + a) / 2);

    // 节点加权
    return dot(m_weight, y);
}

// 一元函数求根
Real FSolve::operator()(Func11 f, Real x)
{
    // Jacobian 算子
    Jacobian j;

    Real dx = 0;
    for (int i = 0; i < ITERATION; i++)
    {
        Real fx = f(x);

        // 获得导数
        Real J = j(f, x);

        // 迭代推进 x = x - f / f'
        dx = fx / J;
        x = x - dx;

        // 当变化量足够小或者是根就退出
        if (fabs(dx) < 1e-8 || fabs(fx) < 1e-8)
            break;
    }

    return x;
}

// 多元向量函数牛顿法求根
void FSolve::operator()(FuncXX f, Vector &x)
{
    // Jacobian 算子
    Jacobian j;

    // 线性求解器
    LinearSolver solver;

    // 创建容器
    Vector dx(x.size());
    for (int i = 0; i < ITERATION; i++)
    {
        // 获得 Jacobi 矩阵
        Matrix J = j(f, x);
        if (J.norm() < ACCURACY)
            break;

        // JT*Jdx = -JT*f
        Vector fx = J.transpose() * f(x);
        Matrix JTJ = J.transpose() * J;

        // 求解最小二乘问题
        dx = solver.solve(JTJ, fx, true);
        x = x - dx;

        // 当变化量足够小或者是根就退出
        if (dx.norm() < ACCURACY || fx.norm() < ACCURACY)
            break;
    }
}